from torch.utils.data import DataLoader

from train.EnergyConfig import EnergyConfig
from train.EnergyDataProvider import StockDataProvider
from train.EnergyDataSet import StockDataSet
from train.EnergyModel import EnergyModel
from train.EnergyTrain import EnergyTrain

if __name__ == '__main__':
    provider = StockDataProvider()
    train_data, test_data, train_value, test_value = provider.get_predicated_data()

    train_dataset = StockDataSet(train_data, train_value)
    test_dataset = StockDataSet(test_data, test_value)

    train_loader = DataLoader(train_dataset, batch_size=EnergyConfig.batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=EnergyConfig.batch_size)

    # init model
    stockModel = EnergyModel().to(EnergyConfig.device)

    # model train
    EnergyTrain(stockModel, train_loader, test_loader).train()
